// NRC语言运行时模块
// 提供手动内存管理等运行时功能

use once_cell::sync::Lazy;
use std::alloc::{alloc, dealloc, Layout};
use std::collections::HashMap;
use std::sync::Mutex;

// 全局内存大小跟踪表
// 使用usize存储指针地址，避免线程安全问题
static MEMORY_SIZE_MAP: Lazy<Mutex<HashMap<usize, usize>>> =
    Lazy::new(|| Mutex::new(HashMap::new()));

/// 引用计数块结构
#[derive(Debug)]
struct RcBlock {
    strong_count: usize, // 强引用计数（非原子版本）
    weak_count: usize,   // 弱引用计数（非原子版本）
    data: *mut u8,       // 实际数据指针
    size: usize,         // 数据大小
}

// 实现Send和Sync，因为我们在单线程环境中使用
unsafe impl Send for RcBlock {}
unsafe impl Sync for RcBlock {}

// 全局引用计数表
static RC_MEMORY_MAP: Lazy<Mutex<HashMap<usize, RcBlock>>> =
    Lazy::new(|| Mutex::new(HashMap::new()));

/// 手动内存管理运行时函数
pub mod memory {
    use super::*;

    /// 分配内存
    /// 这个函数会被LLVM后端调用
    #[no_mangle]
    pub extern "C" fn __malloc(size: usize) -> *mut u8 {
        unsafe {
            let layout = Layout::from_size_align(size, 8).unwrap_or_else(|_| {
                panic!("Invalid layout for malloc: size={}", size);
            });

            let ptr = alloc(layout);
            if ptr.is_null() {
                panic!("Failed to allocate memory of size {}", size);
            }

            // 跟踪内存大小
            if let Ok(mut map) = MEMORY_SIZE_MAP.lock() {
                map.insert(ptr as usize, size);
            }

            // println!("Memory: Allocated {} bytes at {:p}", size, ptr);
            ptr
        }
    }

    /// 释放内存
    /// 这个函数会被LLVM后端调用
    #[no_mangle]
    pub extern "C" fn __free(ptr: *mut u8, size: usize) {
        if ptr.is_null() {
            return;
        }

        unsafe {
            let layout = Layout::from_size_align(size, 8).unwrap_or_else(|_| {
                panic!("Invalid layout for free: size={}", size);
            });

            // 从跟踪表中移除
            if let Ok(mut map) = MEMORY_SIZE_MAP.lock() {
                map.remove(&(ptr as usize));
            }

            // println!("Memory: Freeing {} bytes at {:p}", size, ptr);
            dealloc(ptr, layout);
        }
    }

    /// 简化的释放内存函数 - 只需要指针，自动查找大小
    /// 这个函数可以被NRC代码直接调用
    #[no_mangle]
    pub extern "C" fn free(ptr: *mut u8) {
        if ptr.is_null() {
            return;
        }

        // 从跟踪表中获取大小
        let size = if let Ok(mut map) = MEMORY_SIZE_MAP.lock() {
            map.remove(&(ptr as usize))
        } else {
            None
        };

        if let Some(size) = size {
            unsafe {
                let layout = Layout::from_size_align(size, 8).unwrap_or_else(|_| {
                    panic!("Invalid layout for free: size={}", size);
                });

                // println!("Memory: Freeing {} bytes at {:p} (via free)", size, ptr);
                dealloc(ptr, layout);
            }
        } else {
            // println!("Warning: Attempted to free untracked pointer {:p}", ptr);
        }
    }

    /// 重新分配内存
    /// 这个函数会被LLVM后端调用
    #[no_mangle]
    pub extern "C" fn __realloc(ptr: *mut u8, old_size: usize, new_size: usize) -> *mut u8 {
        if ptr.is_null() {
            return __malloc(new_size);
        }

        if new_size == 0 {
            __free(ptr, old_size);
            return std::ptr::null_mut();
        }

        unsafe {
            let new_layout = Layout::from_size_align(new_size, 8).unwrap_or_else(|_| {
                panic!("Invalid layout for realloc: size={}", new_size);
            });

            let new_ptr = alloc(new_layout);
            if new_ptr.is_null() {
                panic!(
                    "Failed to reallocate memory from {} to {} bytes",
                    old_size, new_size
                );
            }

            // 复制数据
            let copy_size = old_size.min(new_size);
            std::ptr::copy_nonoverlapping(ptr, new_ptr, copy_size);

            // 更新内存跟踪表
            if let Ok(mut map) = MEMORY_SIZE_MAP.lock() {
                map.remove(&(ptr as usize));
                map.insert(new_ptr as usize, new_size);
            }

            // 释放旧内存
            let old_layout = Layout::from_size_align(old_size, 8).unwrap_or_else(|_| {
                panic!("Invalid layout for realloc free: size={}", old_size);
            });
            dealloc(ptr, old_layout);

            // println!("Memory: Reallocated from {} to {} bytes, ptr: {:p} -> {:p}",
            //          old_size, new_size, ptr, new_ptr);
            new_ptr
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_malloc_free() {
        // 测试基本的内存分配和释放
        let size = 1024;
        let ptr = memory::__malloc(size);
        assert!(!ptr.is_null());

        // 写入一些数据
        unsafe {
            std::ptr::write_bytes(ptr, 0x42, size);
        }

        // 释放内存
        memory::__free(ptr, size);
    }

    #[test]
    fn test_realloc() {
        // 测试内存重新分配
        let old_size = 512;
        let new_size = 1024;

        let ptr = memory::__malloc(old_size);
        assert!(!ptr.is_null());

        // 写入一些数据
        unsafe {
            std::ptr::write_bytes(ptr, 0x42, old_size);
        }

        // 重新分配
        let new_ptr = memory::__realloc(ptr, old_size, new_size);
        assert!(!new_ptr.is_null());

        // 释放内存
        memory::__free(new_ptr, new_size);
    }

    #[test]
    fn test_simple_free() {
        // 测试简化的free函数
        let size = 256;
        let ptr = memory::__malloc(size);
        assert!(!ptr.is_null());

        // 写入一些数据
        unsafe {
            std::ptr::write_bytes(ptr, 0xAB, size);
        }

        // 使用简化的free函数释放内存
        memory::free(ptr);
    }

    #[test]
    fn test_free_with_defer_pattern() {
        // 测试free函数与defer模式的配合使用
        let size = 128;
        let ptr = memory::__malloc(size);
        assert!(!ptr.is_null());

        // 模拟defer模式：在函数开始时设置defer，在函数结束时执行
        let mut defer_actions: Vec<Box<dyn Fn()>> = Vec::new();

        // 添加defer动作
        defer_actions.push(Box::new(move || {
            memory::free(ptr);
        }));

        // 使用内存
        unsafe {
            std::ptr::write_bytes(ptr, 0xCD, size);
        }

        // 模拟函数结束，执行defer动作
        for action in defer_actions {
            action();
        }
    }
}

/// 引用计数运行时函数
pub mod rc {
    use super::*;

    /// 创建引用计数对象
    /// 分配内存并初始化引用计数块
    #[no_mangle]
    pub extern "C" fn rc_new(data_size: usize) -> *mut u8 {
        unsafe {
            // 分配数据内存
            let data_layout = Layout::from_size_align(data_size, 8).unwrap_or_else(|_| {
                panic!("Invalid layout for rc_new: size={}", data_size);
            });

            let data_ptr = alloc(data_layout);
            if data_ptr.is_null() {
                panic!("Failed to allocate memory for rc_new: size={}", data_size);
            }

            // 创建引用计数块
            let rc_block = RcBlock {
                strong_count: 1, // 初始强引用计数为1
                weak_count: 0,   // 初始弱引用计数为0
                data: data_ptr,
                size: data_size,
            };

            // 将引用计数块存储到全局表中
            if let Ok(mut map) = RC_MEMORY_MAP.lock() {
                map.insert(data_ptr as usize, rc_block);
            }

            // 同时更新普通内存跟踪表
            if let Ok(mut map) = MEMORY_SIZE_MAP.lock() {
                map.insert(data_ptr as usize, data_size);
            }

            data_ptr
        }
    }

    /// 增加强引用计数
    #[no_mangle]
    pub extern "C" fn rc_increment(ptr: *mut u8) -> *mut u8 {
        if ptr.is_null() {
            return ptr;
        }

        if let Ok(mut map) = RC_MEMORY_MAP.lock() {
            if let Some(block) = map.get_mut(&(ptr as usize)) {
                block.strong_count += 1;
            }
        }
        ptr
    }

    /// 减少强引用计数
    #[no_mangle]
    pub extern "C" fn rc_decrement(ptr: *mut u8) {
        if ptr.is_null() {
            return;
        }

        if let Ok(mut map) = RC_MEMORY_MAP.lock() {
            if let Some(block) = map.get_mut(&(ptr as usize)) {
                block.strong_count -= 1;
                if block.strong_count == 0 {
                    // 强引用计数归零，释放数据
                    if block.weak_count == 0 {
                        // 弱引用计数也为零，完全释放
                        let size = block.size;
                        map.remove(&(ptr as usize));
                        drop(map); // 释放锁

                        if let Ok(mut size_map) = MEMORY_SIZE_MAP.lock() {
                            size_map.remove(&(ptr as usize));
                        }
                        unsafe {
                            let layout = Layout::from_size_align(size, 8).unwrap();
                            dealloc(ptr, layout);
                        }
                    } else {
                        // 还有弱引用，只释放数据，保留控制块
                        unsafe {
                            let layout = Layout::from_size_align(block.size, 8).unwrap();
                            dealloc(block.data, layout);
                        }
                        block.data = std::ptr::null_mut();
                    }
                }
            }
        }
    }

    /// 增加弱引用计数
    #[no_mangle]
    pub extern "C" fn weak_increment(ptr: *mut u8) -> *mut u8 {
        if ptr.is_null() {
            return ptr;
        }

        if let Ok(mut map) = RC_MEMORY_MAP.lock() {
            if let Some(block) = map.get_mut(&(ptr as usize)) {
                block.weak_count += 1;
            }
        }
        ptr
    }

    /// 减少弱引用计数
    #[no_mangle]
    pub extern "C" fn weak_decrement(ptr: *mut u8) {
        if ptr.is_null() {
            return;
        }

        if let Ok(mut map) = RC_MEMORY_MAP.lock() {
            if let Some(block) = map.get_mut(&(ptr as usize)) {
                block.weak_count -= 1;
                if block.weak_count == 0 && block.strong_count == 0 {
                    // 强引用和弱引用都为零，释放控制块
                    map.remove(&(ptr as usize));
                }
            }
        }
    }

    /// 弱引用升级为强引用
    #[no_mangle]
    pub extern "C" fn weak_lock(ptr: *mut u8) -> *mut u8 {
        if ptr.is_null() {
            return std::ptr::null_mut();
        }

        if let Ok(mut map) = RC_MEMORY_MAP.lock() {
            if let Some(block) = map.get_mut(&(ptr as usize)) {
                if block.strong_count > 0 && !block.data.is_null() {
                    block.strong_count += 1;
                    return block.data;
                }
            }
        }
        std::ptr::null_mut() // 升级失败
    }

    /// 检查引用计数对象是否有效
    #[no_mangle]
    pub extern "C" fn rc_is_valid(ptr: *mut u8) -> bool {
        if ptr.is_null() {
            return false;
        }

        if let Ok(map) = RC_MEMORY_MAP.lock() {
            if let Some(block) = map.get(&(ptr as usize)) {
                return block.strong_count > 0 && !block.data.is_null();
            }
        }
        false
    }

    /// 获取强引用计数（调试用）
    #[no_mangle]
    pub extern "C" fn rc_strong_count(ptr: *mut u8) -> usize {
        if ptr.is_null() {
            return 0;
        }

        if let Ok(map) = RC_MEMORY_MAP.lock() {
            if let Some(block) = map.get(&(ptr as usize)) {
                return block.strong_count;
            }
        }
        0
    }

    /// 获取弱引用计数（调试用）
    #[no_mangle]
    pub extern "C" fn rc_weak_count(ptr: *mut u8) -> usize {
        if ptr.is_null() {
            return 0;
        }

        if let Ok(map) = RC_MEMORY_MAP.lock() {
            if let Some(block) = map.get(&(ptr as usize)) {
                return block.weak_count;
            }
        }
        0
    }
}
